package util;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.CtNewMethod;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

/**
 * 方式1
 * 检测方法的执行时间
 * 会出现如下问题：
 * 当修改扣发代码如下时，因代码是不合理的，会报错：Exception in thread "main" java.lang.VerifyError: Expecting a stack map frame
 * public void run(){
 *     long start = System.currentTimeMillis();
 *     int i = 1;
 *     while(true){
 *         i = print(i);
 *     }
 *     System.out.println("run() cost: "+ (System.currentTimeMillis() - start));
 * }
 * @author: Alan
 * @date: 2021/4/25 22:41
 */
public class CostTime1 implements ClassFileTransformer {


    public CostTime1() {
        System.out.println("CostTime1");
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        // 这里我们限制下，只针对目标包下进行耗时统计
        if (!className.startsWith("com/agent/demo/TimeTest")) {
            return classfileBuffer;
        }

        CtClass ctclass = null;
        try {

            ClassPool pool = ClassPool.getDefault();
            ctclass = pool.makeClass(new ByteArrayInputStream(classfileBuffer));

            for (CtMethod method : ctclass.getDeclaredMethods()) {
                // 所有方法，统计耗时；请注意，需要通过`addLocalVariable`来声明局部变量
                method.addLocalVariable("startTimeAgent", CtClass.longType);
                method.insertBefore("startTimeAgent = System.currentTimeMillis();");
                String methodName = method.getLongName();
                method.insertAfter("System.out.println(\"" + methodName +
                        " cost1: \" + (System.currentTimeMillis() - startTimeAgent));");
            }
            byte[] transformed = ctclass.toBytecode();
            return transformed;
        } catch (Exception e) {
            System.out.println(e.getMessage());
            e.printStackTrace();
        }

        return null;
    }
}
